# Copyright 2013-present Barefoot Networks, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Thrift PD interface basic tests
"""

from collections import OrderedDict

import time
import sys
import logging
import copy
import pdb

import unittest
import random

import pd_base_tests

from ptf import config
from ptf.testutils import *
from ptf.thriftutils import *

import os

from pal_rpc.ttypes import *
from pcie_pkt_test.p4_pd_rpc.ttypes import *
from pkt_pd_rpc.ttypes import *
from mc_pd_rpc.ttypes import *
from res_pd_rpc.ttypes import *
from devport_mgr_pd_rpc.ttypes import *

dev_id = 0
MAX_PORT_COUNT = 456

global sess_pkt0
global sess_pkt1
global sess_pkt2
global sess_pkt3

# First grab all ports
all_ports = []
swports = []
for device, port, ifname in config["interfaces"]:
  all_ports.append(port)
  all_ports.sort()
# Default to a few ports if we didn't find any in the config
if all_ports == []:
  all_ports = range(9)
# Next, filter the ports based on the allowed pipes
num_pipes = int(test_param_get('num_pipes'))
all_pipes = range(num_pipes)
for p in all_ports:
  pipe = p >> 7
  if pipe in all_pipes:
    swports.append(p)
print "All ports", swports

def port_tbl_add(test, shdl, dev_tgt, port, dport):
    match_spec = pcie_pkt_test_port_tbl_match_spec_t( hex_to_i16(port) )
    action_spec = pcie_pkt_test_set_md_action_spec_t( hex_to_i16(dport) )
    entry_hdl = test.client.port_tbl_table_add_with_set_md( shdl, dev_tgt, match_spec, action_spec )
    return entry_hdl

def match_pkt_string(test, to_pkt, str1):
    if (str(to_pkt) == str(str1)):
        logging.debug("matched")
        return True
    else:
        logging.debug("not matched")
        return False

class TestPG(pd_base_tests.ThriftInterfaceDataPlane):
    def __init__(self):
        pd_base_tests.ThriftInterfaceDataPlane.__init__(self, ["pcie_pkt_test"])

    def verify_one_pass(self):
        sess_pkt0 = sess_pkt1 = sess_pkt2 = sess_pkt3 = -1;
        try:
            self.pkt.init()
            sess_pkt0 = self.pkt.client_init()
            sess_pkt1 = self.pkt.client_init()
            sess_pkt2 = self.pkt.client_init()
            sess_pkt3 = self.pkt.client_init()
            p0 = simple_eth_packet(pktlen=1000)
            p1 = simple_eth_packet(pktlen=100)
            p2 = simple_eth_packet(pktlen=500)
            p3 = simple_eth_packet(pktlen=1400)
            self.pkt.test_send_pkt(sess_pkt0, str(p0), len(p0), 0)
            self.pkt.test_send_pkt(sess_pkt1, str(p1), len(p1), 1)
            self.pkt.test_send_pkt(sess_pkt2, str(p2), len(p2), 2)
            self.pkt.test_send_pkt(sess_pkt3, str(p3), len(p3), 3)
            print "Waiting for DMA to complete"
            sys.stdout.flush()
            time.sleep(5)
            self.pkt.test_verify_pkt_tx(sess_pkt0)
            self.pkt.test_verify_pkt_tx(sess_pkt1)
            self.pkt.test_verify_pkt_tx(sess_pkt2)
            self.pkt.test_verify_pkt_tx(sess_pkt3)
            verify_packet(self, p0, swports[0])
            verify_packet(self, p1, swports[0])
            verify_packet(self, p2, swports[0])
            verify_packet(self, p3, swports[0])
            verify_no_other_packets(self)


            p = simple_eth_packet(pktlen=256, eth_dst='80:01:02:03:04:05')
            send_packet(self, swports[0], str(p))
            time.sleep(1)
            rx_pkt = self.pkt.test_get_pkt_rx(sess_pkt0)
            retValue = match_pkt_string(self, p, rx_pkt.buf)
            self.assertTrue(retValue == True, "Rx pkt did not match the Tx1pkt")

            p = simple_eth_packet(pktlen=708, eth_dst='81:01:02:03:04:05')
            send_packet(self, swports[0], str(p))
            time.sleep(1)
            rx_pkt = self.pkt.test_get_pkt_rx(sess_pkt0)
            retValue = match_pkt_string(self, p, rx_pkt.buf)
            self.assertTrue(retValue == True, "Rx pkt did not match the Tx2pkt")

            p = simple_eth_packet(pktlen=1200, eth_dst='85:01:02:03:04:05')
            send_packet(self, swports[0], str(p))
            time.sleep(1)
            rx_pkt = self.pkt.test_get_pkt_rx(sess_pkt0)
            retValue = match_pkt_string(self, p, rx_pkt.buf)
            self.assertTrue(retValue == True, "Rx pkt did not match the Tx3pkt")

            p = simple_eth_packet(pktlen=500, eth_dst='82:01:02:03:04:05')
            send_packet(self, swports[0], str(p))
            time.sleep(1)
            rx_pkt = self.pkt.test_get_pkt_rx(sess_pkt0)
            retValue = match_pkt_string(self, p, rx_pkt.buf)
            self.assertTrue(retValue == True, "Rx pkt did not match the Tx4pkt")

            p = simple_eth_packet(pktlen=873, eth_dst='87:01:02:03:04:05')
            send_packet(self, swports[0], str(p))
            time.sleep(1)
            rx_pkt = self.pkt.test_get_pkt_rx(sess_pkt0)
            retValue = match_pkt_string(self, p, rx_pkt.buf)
            self.assertTrue(retValue == True, "Rx pkt did not match the Tx5pkt")

            p = simple_eth_packet(pktlen=1105, eth_dst='84:01:02:03:04:05')
            send_packet(self, swports[0], str(p))
            time.sleep(1)
            rx_pkt = self.pkt.test_get_pkt_rx(sess_pkt0)
            retValue = match_pkt_string(self, p, rx_pkt.buf)
            self.assertTrue(retValue == True, "Rx pkt did not match the Tx6pkt")

            p = simple_eth_packet(pktlen=302, eth_dst='86:01:02:03:04:05')
            send_packet(self, swports[0], str(p))
            time.sleep(1)
            rx_pkt = self.pkt.test_get_pkt_rx(sess_pkt0)
            retValue = match_pkt_string(self, p, rx_pkt.buf)
            self.assertTrue(retValue == True, "Rx pkt did not match the Tx7pkt")

            p = simple_eth_packet(pktlen=83, eth_dst='83:01:02:03:04:05')
            send_packet(self, swports[0], str(p))
            time.sleep(1)
            rx_pkt = self.pkt.test_get_pkt_rx(sess_pkt0)
            retValue = match_pkt_string(self, p, rx_pkt.buf)
            self.assertTrue(retValue == True, "Rx pkt did not match the Tx8pkt")

            p = simple_eth_packet(pktlen=803, eth_dst='83:01:02:03:04:05')
            send_packet(self, swports[0], str(p))
            time.sleep(1)
            rx_pkt = self.pkt.test_get_pkt_rx(sess_pkt0)
            retValue = match_pkt_string(self, p, rx_pkt.buf)
            self.assertTrue(retValue == True, "Rx pkt did not match the Tx9pkt")

        finally:
            if (sess_pkt3 != -1):
                self.pkt.client_cleanup(sess_pkt3)
            if (sess_pkt0 != -1):
                self.pkt.client_cleanup(sess_pkt0)
            if (sess_pkt2 != -1):
                self.pkt.client_cleanup(sess_pkt2)
            if (sess_pkt1 != -1):
                self.pkt.client_cleanup(sess_pkt1)

            self.pkt.cleanup()

    def configure(self, sess_hdl, dev_id, port_tbl_handles):
        num_pipes = int(test_param_get('num_pipes'))

        dev_tgt = DevTarget_t(dev_id, hex_to_i16(0xFFFF))
        dev_tgts = [DevTarget_t(dev_id, hex_to_i16(0)),
                    DevTarget_t(dev_id, hex_to_i16(1)),
                    DevTarget_t(dev_id, hex_to_i16(2)),
                    DevTarget_t(dev_id, hex_to_i16(3))]

        prop = tbl_property_t.TBL_PROP_TBL_ENTRY_SCOPE
        prop_val = tbl_property_value_t.ENTRY_SCOPE_SINGLE_PIPELINE
        self.client.port_tbl_set_property(sess_hdl, dev_id, prop, prop_val, 0)

        # Program tables such that a packet coming in on port 0 will go to port
        # 320, the CPU PCIe port.  Note that on two pipe SKUs the PCIe port
        # moves to pipe 1 port 64 which is 192.
        if num_pipes > 2:
            cpu_pcie_port = 320
        else:
            cpu_pcie_port = 192
        entry_hdl = port_tbl_add(self, sess_hdl, dev_tgts[swports[0] >> 7], swports[0], cpu_pcie_port)
        port_tbl_handles.append(entry_hdl)

        # Program tables such that a packet coming in on port 320 (CPU PCIe)
        # will go to port 1.
        entry_hdl = port_tbl_add(self, sess_hdl, dev_tgts[cpu_pcie_port >> 7], cpu_pcie_port, swports[0])
        port_tbl_handles.append(entry_hdl)

    def addPorts(self):
        for i in swports:
            self.pal.pal_port_add(dev_id, i, pal_port_speed_t.BF_SPEED_10G, pal_fec_type_t.BF_FEC_TYP_NONE)
            self.pal.pal_port_an_set(dev_id, i, pal_autoneg_policy_t.BF_AN_FORCE_DISABLE)
            self.pal.pal_port_enable(dev_id, i)
        ports_not_up = True
        print "Waiting for ports to come up..."
        sys.stdout.flush()
        num_tries = 12
        i = 0
        while ports_not_up:
            ports_not_up = False
            for p in swports:
                x = self.pal.pal_port_oper_status_get(dev_id, p)
                if x == pal_oper_status_t.BF_PORT_DOWN:
                    ports_not_up = True
                    print "  port", p, "is down"
                    sys.stdout.flush()
                    time.sleep(5)
                    break
            i += 1
            if i >= num_tries:
                break
        assert ports_not_up == False

    def runTest(self):
        num_pipes = int(test_param_get('num_pipes'))
        print "Num pipes is", num_pipes
        sys.stdout.flush()

        port_tbl_handles = []

        try:
            sess_hdl = self.conn_mgr.client_init()
            self.configure(sess_hdl, dev_id, port_tbl_handles)

            # Wait for all pipe APIs to complete.
            self.conn_mgr.complete_operations(sess_hdl)

            print 'Testing after cold boot'
            sys.stdout.flush()
            self.verify_one_pass()
            print '  Done'
            sys.stdout.flush()

            print 'Initiating fast reconfig'
            sys.stdout.flush()
            #Perform a fast reconfig
            port_tbl_handles = []
            self.devport_mgr.devport_mgr_warm_init_begin(dev_id, dev_init_mode.DEV_WARM_INIT_FAST_RECFG, dev_serdes_upgrade_mode.DEV_SERDES_UPD_NONE, True)

            if test_param_get('target') == 'hw':
                self.addPorts()

            print 'Testing during warm init'
            sys.stdout.flush()
            self.verify_one_pass()
            print '  Done'
            sys.stdout.flush()

            self.configure(sess_hdl, dev_id, port_tbl_handles)
            self.devport_mgr.devport_mgr_warm_init_end(dev_id)

            self.conn_mgr.complete_operations(sess_hdl)
            print 'Testing after warm init completed'
            sys.stdout.flush()
            self.verify_one_pass()
            print '  Done'
            sys.stdout.flush()

            print 'Success!'
            sys.stdout.flush()

        finally:
            print 'Cleaning up'
            sys.stdout.flush()
            for eh in port_tbl_handles:
                sts = self.client.port_tbl_table_delete(sess_hdl, dev_id, eh)

            sess_hdl = self.conn_mgr.client_cleanup(sess_hdl)

            print 'Cleanup done'
            sys.stdout.flush()
